package util

import (
	"bytes"
	"encoding/base64"
	"gitee.com/ctyunfaas/cf-runtime-go-sdk/config"
	"log"
	"net/http"
	"os"
)

func GetEnvOrDefault(key, defaultValue string) string {
	value, exists := os.LookupEnv(key)
	if exists {
		return value
	}
	return defaultValue
}

// LoggerInterceptor 拦截日志并同时输出到标准输出和内部缓冲区
type LoggerInterceptor struct {
	LogBuffer bytes.Buffer
}

// 实现 io.Write 接口
func (li *LoggerInterceptor) Write(p []byte) (n int, err error) {
	// 写入标准输出
	if _, err := os.Stdout.Write(p); err != nil {
		return 0, err // 返回写入错误
	}

	// 检查写入后的缓冲区大小
	if li.LogBuffer.Len()+len(p) > config.MAX_TAIL_LOG_BYTES {
		// 如果超过最大大小，则移除旧的日志以适应新的日志
		extra := (li.LogBuffer.Len() + len(p)) - config.MAX_TAIL_LOG_BYTES
		li.LogBuffer.Next(extra)
	}

	// 写入内部缓冲区
	return li.LogBuffer.Write(p)
}

func NewLoggerInterceptor() *LoggerInterceptor {
	return &LoggerInterceptor{}
}

// 包装 http.ResponseWriter, 实现统一返回 tail log
type customResponseWriter struct {
	http.ResponseWriter
	useTail bool
	*LoggerInterceptor
	bodyBuffer       bytes.Buffer
	statusCodeBuffer int
}

func (w *customResponseWriter) FinishResponse() {
	if w.useTail {
		w.addCustomHeader()
	}
	// 写入状态码
	w.ResponseWriter.WriteHeader(w.statusCodeBuffer)
	// 写入响应体
	if _, err := w.ResponseWriter.Write(w.bodyBuffer.Bytes()); err != nil {
		log.Printf("Error writing body: %v", err)
	}
}

func (w *customResponseWriter) WriteHeader(code int) {
	w.statusCodeBuffer = code
}

func (w *customResponseWriter) Write(b []byte) (int, error) {
	return w.bodyBuffer.Write(b)
}

func (w *customResponseWriter) addCustomHeader() {
	// 只设置一次自定义 header
	if w.useTail {
		w.useTail = false
		// 获取LogBuffer的字符串表示
		logStr := w.LoggerInterceptor.LogBuffer.String()
		// 使用base64.StdEncoding.EncodeToString方法将字符串进行base64编码
		encodedLogStr := base64.StdEncoding.EncodeToString([]byte(logStr))
		w.Header().Set(config.LOG_RESULT, encodedLogStr)
	}
}

func NewCustomResponseWriter(w http.ResponseWriter, useTail bool, li *LoggerInterceptor) *customResponseWriter {
	return &customResponseWriter{w, useTail, li, bytes.Buffer{}, http.StatusOK}
}
